from flask import Flask, render_template, request, jsonify
import os
from utils.yolo import Helmet_Detector
import uuid
app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'static/uploads'
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg', 'gif'}

# 初始化YOLO模型
yolo = Helmet_Detector()


# 检查文件扩展名是否允许
def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']


@app.route('/')
def index():
    return render_template('index.html')


@app.route('/detect', methods=['POST'])
def detect():
    if 'file' not in request.files:
        return jsonify({'error': 'No file part'})

    file = request.files['file']

    if file.filename == '':
        return jsonify({'error': 'No selected file'})
    if file and allowed_file(file.filename):
        random_filename = f"{uuid.uuid4().hex}{os.path.splitext(file.filename)[1]}"
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], random_filename)
        file.save(filepath)
        filepath = filepath.replace('\\', '/')
        # 使用YOLO模型进行检测
        try:
            results = yolo.detect(filepath)
            return jsonify({'success': True, 'results': results, 'image_path': filepath})
        except Exception as e:
            print(e)
            return jsonify({'error': str(e)})

    return jsonify({'error': 'Invalid file format'})


if __name__ == '__main__':
    # 确保上传文件夹存在
    os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
    app.run(debug=True)